__global__ void tecoKernelArgmaxFT16Algo76_ziyang_(ArgMaxArgs arg) {
    const int high_num = arg.high_num;
    const int axis_num = arg.axis_num;
    const int low_num = arg.low_num;
    // 张量的一些尺寸数据
    const half *x = (const half *)arg.x;
    int64_t *y = (int64_t *)arg.y;
    // 输入数据指针和输出索引的指针
    const int spe_num = arg.spe_num;  // SPE的数量
    Stride stride;                    // 步幅？？？

    // deal_high_num_per_spe 处理高维数据的每一个SPE
    int deal_high_num_per_spe = high_num / spe_num;
    int deal_high_num_per_spe_rem = high_num % spe_num;  // 均分后剩余的任务数量
    // 数据的均匀划分
    if (threadIdx < deal_high_num_per_spe_rem) {
        ++deal_high_num_per_spe;
    }
    // 【以及多余数据的分配】

    int axis_time = ROUND_UP(axis_num, SPM_MAX_NUM);
    int axis_rem = axis_num % SPM_MAX_NUM;
    if (deal_high_num_per_spe == 0) return;
    half *spm_x = (half *)rt_spm_malloc(SPM_MAX_BYTE);
    int init_index = threadIdx < deal_high_num_per_spe_rem ?
                         threadIdx * deal_high_num_per_spe :
                         threadIdx * deal_high_num_per_spe + deal_high_num_per_spe_rem;
    // #pragma omp parallel for
    // 没有omp环境
    for (int i = 0; i < deal_high_num_per_spe; i++) {
        int offset_h = (init_index + i) * axis_num * low_num;

        for (int l = 0; l < low_num; l++) {
            int index = offset_h + l;
            int offset_l = offset_h + l;
            int index_max = 0;
            int deal_num = 0;
            int offset_a = 0;
            half value_max = FLT16_MIN_VALUE;

            for (int at = 0; at < axis_time; at++) {
                deal_num = at == axis_time - 1 ? axis_rem : SPM_MAX_NUM;
                offset_a = offset_l + at * SPM_MAX_NUM * low_num;
                stride_config(stride, deal_num, (low_num - 1) * sizeof(half));
                memcpy_stride(spm_x, x + offset_a, sizeof(half), stride);

                for (int d = 0; d < deal_num; d++) {
                    if (value_max < spm_x[d]) {
                        value_max = spm_x[d];
                        index_max = at * SPM_MAX_NUM + d;
                    }
                }
            }
            *(y + (init_index + i) * low_num + l) = index_max;
        }
    }
    rt_spm_free(spm_x);
    return;
}

// 这里写自己的kernel函数--注意，函数名一定不要相同！！！
__global__ void tecoKernelArgmaxFT16Algo76_yourname1_(ArgMaxArgs arg) {
    const int high_num = arg.high_num;
    const int axis_num = arg.axis_num;
    const int low_num = arg.low_num;
    const half *x = (const half *)arg.x;
    int64_t *y = (int64_t *)arg.y;
    const int spe_num = arg.spe_num;
    Stride stride;

    int deal_high_num_per_spe = high_num / spe_num;
    int deal_high_num_per_spe_rem = high_num % spe_num;
    if (threadIdx < deal_high_num_per_spe_rem) {
        ++deal_high_num_per_spe;
    }
    int axis_time = ROUND_UP(axis_num, SPM_MAX_NUM);
    int axis_rem = axis_num % SPM_MAX_NUM;
    if (deal_high_num_per_spe == 0) return;
    half *spm_x = (half *)rt_spm_malloc(SPM_MAX_BYTE);
    int init_index = threadIdx < deal_high_num_per_spe_rem ?
                         threadIdx * deal_high_num_per_spe :
                         threadIdx * deal_high_num_per_spe + deal_high_num_per_spe_rem;

    for (int i = 0; i < deal_high_num_per_spe; i++) {
        int offset_h = (init_index + i) * axis_num * low_num;
        for (int l = 0; l < low_num; l++) {
            int index = offset_h + l;
            int offset_l = offset_h + l;
            int index_max = 0;
            int deal_num = 0;
            int offset_a = 0;
            half value_max = FLT16_MIN_VALUE;
            for (int at = 0; at < axis_time; at++) {
                deal_num = at == axis_time - 1 ? axis_rem : SPM_MAX_NUM;
                offset_a = offset_l + at * SPM_MAX_NUM * low_num;
                stride_config(stride, deal_num, (low_num - 1) * sizeof(half));
                memcpy_stride(spm_x, x + offset_a, sizeof(half), stride);
                for (int d = 0; d < deal_num; d++) {
                    if (value_max < spm_x[d]) {
                        value_max = spm_x[d];
                        index_max = at * SPM_MAX_NUM + d;
                    }
                }
            }
            *(y + (init_index + i) * low_num + l) = index_max;
        }
    }
    rt_spm_free(spm_x);
    return;
}

// 这里写自己的kernel函数--注意，函数名一定不要相同！！！
__global__ void tecoKernelArgmaxFT16Algo76_yourname2_(ArgMaxArgs arg) {
    const int high_num = arg.high_num;
    const int axis_num = arg.axis_num;
    const int low_num = arg.low_num;
    const half *x = (const half *)arg.x;
    int64_t *y = (int64_t *)arg.y;
    const int spe_num = arg.spe_num;
    Stride stride;

    int deal_high_num_per_spe = high_num / spe_num;
    int deal_high_num_per_spe_rem = high_num % spe_num;
    if (threadIdx < deal_high_num_per_spe_rem) {
        ++deal_high_num_per_spe;
    }
    int axis_time = ROUND_UP(axis_num, SPM_MAX_NUM);
    int axis_rem = axis_num % SPM_MAX_NUM;
    if (deal_high_num_per_spe == 0) return;
    half *spm_x = (half *)rt_spm_malloc(SPM_MAX_BYTE);
    int init_index = threadIdx < deal_high_num_per_spe_rem ?
                         threadIdx * deal_high_num_per_spe :
                         threadIdx * deal_high_num_per_spe + deal_high_num_per_spe_rem;

    for (int i = 0; i < deal_high_num_per_spe; i++) {
        int offset_h = (init_index + i) * axis_num * low_num;
        for (int l = 0; l < low_num; l++) {
            int index = offset_h + l;
            int offset_l = offset_h + l;
            int index_max = 0;
            int deal_num = 0;
            int offset_a = 0;
            half value_max = FLT16_MIN_VALUE;
            for (int at = 0; at < axis_time; at++) {
                deal_num = at == axis_time - 1 ? axis_rem : SPM_MAX_NUM;
                offset_a = offset_l + at * SPM_MAX_NUM * low_num;
                stride_config(stride, deal_num, (low_num - 1) * sizeof(half));
                memcpy_stride(spm_x, x + offset_a, sizeof(half), stride);

                for (int d = 0; d < deal_num; d++) {
                    if (value_max < spm_x[d]) {
                        value_max = spm_x[d];
                        index_max = at * SPM_MAX_NUM + d;
                    }
                }
            }
            *(y + (init_index + i) * low_num + l) = index_max;
        }
    }
    rt_spm_free(spm_x);
    return;
}

__global__ void tecoKernelArgmaxFT16Algo76_yuanma_(ArgMaxArgs arg) {
    const int high_num = arg.high_num;
    const int axis_num = arg.axis_num;
    const int low_num = arg.low_num;
    // 张量的一些尺寸数据
    const half *x = (const half *)arg.x;
    int64_t *y = (int64_t *)arg.y;
    // 输入数据指针和输出索引的指针
    const int spe_num = arg.spe_num;  // SPE的数量
    Stride stride;                    // 步幅？？？

    // deal_high_num_per_spe 处理高维数据的每一个SPE
    int deal_high_num_per_spe = high_num / spe_num;
    int deal_high_num_per_spe_rem = high_num % spe_num;  // 均分后剩余的任务数量
    // 数据的均匀划分
    if (threadIdx < deal_high_num_per_spe_rem) {
        ++deal_high_num_per_spe;
    }
    // 【以及多余数据的分配】

    int axis_time = ROUND_UP(axis_num, SPM_MAX_NUM);
    int axis_rem = axis_num % SPM_MAX_NUM;
    if (deal_high_num_per_spe == 0) return;
    half *spm_x = (half *)rt_spm_malloc(SPM_MAX_BYTE);
    int init_index = threadIdx < deal_high_num_per_spe_rem ?
                         threadIdx * deal_high_num_per_spe :
                         threadIdx * deal_high_num_per_spe + deal_high_num_per_spe_rem;
    // #pragma omp parallel for
    // 没有omp环境
    for (int i = 0; i < deal_high_num_per_spe; i++) {
        int offset_h = (init_index + i) * axis_num * low_num;

        for (int l = 0; l < low_num; l++) {
            int index = offset_h + l;
            int offset_l = offset_h + l;
            int index_max = 0;
            int deal_num = 0;
            int offset_a = 0;
            half value_max = FLT16_MIN_VALUE;

            for (int at = 0; at < axis_time; at++) {
                deal_num = at == axis_time - 1 ? axis_rem : SPM_MAX_NUM;
                offset_a = offset_l + at * SPM_MAX_NUM * low_num;
                stride_config(stride, deal_num, (low_num - 1) * sizeof(half));
                memcpy_stride(spm_x, x + offset_a, sizeof(half), stride);

                for (int d = 0; d < deal_num; d++) {
                    if (value_max < spm_x[d]) {
                        value_max = spm_x[d];
                        index_max = at * SPM_MAX_NUM + d;
                    }
                }
            }
            *(y + (init_index + i) * low_num + l) = index_max;
        }
    }
    rt_spm_free(spm_x);
    return;
}
